import asyncio
import logging
import sys
import time
import warnings
from typing import List

import aiomysql
import pymysql

warnings.filterwarnings('ignore', category=aiomysql.Warning)


class ColorFormatter(logging.Formatter):
    BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)

    COLORS = {
        'WARNING': MAGENTA,
        'INFO': BLUE,
        'DEBUG': WHITE,
        'CRITICAL': RED,
        'ERROR': RED,
        'RED': RED,
        'GREEN': GREEN,
        'YELLOW': YELLOW,
        'BLUE': BLUE,
        'MAGENTA': MAGENTA,
        'CYAN': CYAN,
        'WHITE': WHITE,
    }

    RESET_SEQ = "\033[0m"
    COLOR_SEQ = "\033[1;%dm"
    BOLD_SEQ = "\033[1m"

    def __init__(self, *args, **kwargs):
        logging.Formatter.__init__(self, *args, **kwargs)

    def format(self, record):
        levelname = record.levelname
        color = ColorFormatter.COLOR_SEQ % (30 + ColorFormatter.COLORS[levelname])
        message = logging.Formatter.format(self, record)
        message = message.replace("$RESET", ColorFormatter.RESET_SEQ) \
            .replace("$BOLD", ColorFormatter.BOLD_SEQ) \
            .replace("$COLOR", color)
        for k, v in ColorFormatter.COLORS.items():
            message = message.replace("$" + k, ColorFormatter.COLOR_SEQ % (v + 30)) \
                .replace("$BG" + k, ColorFormatter.COLOR_SEQ % (v + 40)) \
                .replace("$BG-" + k, ColorFormatter.COLOR_SEQ % (v + 40))
        return message + ColorFormatter.RESET_SEQ


console_log_format = '%(asctime)s [$COLOR%(levelname)s$RESET] %(filename)s:%(lineno)s [%(threadName)s] - $COLOR%(message)s$RESET'
stream_formatter = ColorFormatter(console_log_format)

stream_handler = logging.StreamHandler(stream=sys.stdout)
stream_handler.setFormatter(stream_formatter)

logger = logging.getLogger()
logger.addHandler(stream_handler)
logger.setLevel('DEBUG')


class DataBase:

    def __init__(self, database_config: dict):
        self._database_config = database_config
        self._opening_connection = []

    async def get_db(self):
        con = await aiomysql.connect(**self._database_config)
        cursor = await con.cursor()
        self._opening_connection.append((con, cursor))
        return con, cursor

    @property
    def database(self):
        return self._database_config.get('db')

    @property
    def user(self):
        return self._database_config.get('user')

    @property
    def host(self):
        return '{}:{}'.format(self._database_config.get('host'), self._database_config.get('port'))


database_error_code = {
    1050: lambda x: x,  # 表已存在
    1215: '由于外键无法创建表',
    1142: '命令无权限',
    1709: "索引限制",
}
PASS_CODE = [1050]
RETRY_CODE = [1215]
ERR_CODE = [1142, 1709]


class SyncDatabase:
    SignalQuit = 'quit'

    def __init__(self,
                 src: DataBase,
                 dst: DataBase,
                 chunk_size=5000,
                 coroutine=5,
                 override: bool = False,
                 include_tables: List[str] = None,
                 exclude_tables: List[str] = None
                 ):
        '''

        :params src: 源数据库
        :params dst: 目标数据库
        :params chunk_size: 一次查询的条目数量
        :params coroutine: 同时运行的协程数量
        :params override: True会删除当前表，重新建表同步，False在元数据的基础上同步新数据
        :params include_tables: 需要同步的数据表
        :params exclude_tables: 不需要同步的数据表
        '''
        self._initial()
        self._src = src
        self._dst = dst
        logger.info('数据源信息:{}-{}'.format(self._src.host, self._src.database))
        logger.info('目标库信息:{}-{}'.format(self._dst.host, self._dst.database))
        self._include_tables = include_tables if include_tables else []
        self._exclude_tables = exclude_tables if exclude_tables else []
        self._sync_tables = []  # 本次实际要同步的表
        self._chunk_size = chunk_size
        self._coroutine = coroutine
        self._override = override
        if all([self._include_tables, self._exclude_tables]):
            raise Exception('include_tables、exclude_tables仅可选其一或者不选')

        self._retry_sql = []  # 创表失败语句
        self._err_sql = []
        self._retry_count = 5

    def _initial(self):
        def _escape_float(value, mapping=None):
            s = repr(value)
            if s in ("inf", "nan"):
                return 'NULL'
            if "e" not in s:
                s += "e0"
            return s

        pymysql.converters._escape_table[ord("%")] = "%%"  # todo json字符串中的%会重复
        pymysql.converters.encoders[float] = _escape_float

        logger.warning('已修改pymysql.converters._escape_table')
        logger.warning('已修改pymysql.converters.encoders')

    async def delete_table(self, database: DataBase):
        if self._override:
            con, cursor = await database.get_db()
            # check = input('当前override为True,表示数据表会覆盖,请确认被同步数据链接是否正确\n'
            #               '当前被同步数据库名信息:{}-{}\n(YES/NO):'.format(database.host, database.database))
            # if check != 'YES':
            #     logger.info('未确认,数据表未删除')
            #     con.close()
            #     return

            for table in self._sync_tables:
                sql = """DROP TABLE IF EXISTS `{}`;""".format(table)
                await cursor.execute(sql)
            con.close()

    async def retry_sync_table_structure(self, cursor: aiomysql.Cursor):
        tmp = []
        for _ in range(self._retry_count):
            if not self._retry_sql:
                break
            for sql in self._retry_sql:
                try:
                    await cursor.execute(sql)
                except:
                    tmp.append(sql)
            self._retry_sql = tmp
            tmp = []
        for sql in tmp:
            self._err_sql.append(sql)

    async def table_inspector(self, cursor: aiomysql.Cursor, table):
        sql = """SHOW CREATE TABLE `{table_name}`""".format(table_name=table)
        try:
            await cursor.execute(sql)
            inspector = await cursor.fetchall()
            return inspector[0][1]
        except pymysql.err.OperationalError as e:
            # 1142 command denied
            logger.warning(','.join([str(i) for i in e.args]))
        except pymysql.err.ProgrammingError as e:
            logger.error(','.join([str(i) for i in e.args]))

    async def database_inspector(self, cursor: aiomysql.Cursor):
        sql = """SHOW TABLES;"""
        try:
            await cursor.execute(sql)
            all_tables = await cursor.fetchall()
            return [t[0] for t in all_tables]
        except Exception as e:
            logger.error(''.join([str(i) for i in e.args]))
            sys.exit(-1)

    async def sync_table_structure(self):
        wait_to_exec = []  # 即将执行的建表语句
        src_con, src_cursor = await self._src.get_db()
        dst_con, dst_cursor = await self._dst.get_db()

        async def _set_one(t):
            t_structure = await self.table_inspector(src_cursor, t)
            if t_structure:
                wait_to_exec.append(t_structure)
                self._sync_tables.append(t)

        if self._include_tables:
            for table in self._include_tables:
                await _set_one(table)
        elif self._exclude_tables:
            all_tables = await self.database_inspector(src_cursor)
            for exclude_table in self._exclude_tables:
                all_tables.remove(exclude_table)
            for table in all_tables:
                await _set_one(table)
        else:
            all_tables = await self.database_inspector(src_cursor)
            for table in all_tables:
                await _set_one(table)
        await self.delete_table(self._dst)
        print(len(wait_to_exec))
        for sql in wait_to_exec:
            try:
                await dst_cursor.execute(sql)
            except Exception as e:
                print(e.args)
                err_code, err_msg = e.args
                if err_code in PASS_CODE:
                    continue
                elif err_code in RETRY_CODE:
                    self._retry_sql.append(sql)
                elif err_code in ERR_CODE:
                    self._err_sql.append(sql)
        await self.retry_sync_table_structure(dst_cursor)
        src_con.close()
        dst_con.close()

    async def _sync_one_table(self, src_con, src_cursor, dst_con, dst_cursor, table):
        t = time.time()
        logger.info('正在同步 「{}」...'.format(table))
        table_labels, labels_count = None, 0

        init_number, finish_number = 0, self._chunk_size
        while 1:
            select_sql = """SELECT * FROM `{table_name}` limit {init},{finish};""".format(
                table_name=table,
                init=init_number,
                finish=finish_number)

            await src_cursor.execute(select_sql)
            src_data = await src_cursor.fetchall()
            if not src_data:
                break

            init_number = init_number + finish_number
            finish_number = finish_number + self._chunk_size

            if not table_labels:
                labels = [item[0] for item in src_cursor.description]
                labels_count = len(labels)
                table_labels = ','.join(['`{}`'.format(item) for item in labels])

            insert_sql = """INSERT IGNORE INTO `{table_name}` ({keys}) VALUES ({values});""".format(
                table_name=table,
                keys=table_labels,
                values=','.join(['%s' for i in range(labels_count)]))
            try:
                await dst_cursor.executemany(insert_sql, src_data)
            except Exception as e:
                logger.error('「{data_base}」-「{table_name}」\n「{exception}」\n「{select_sql}」'.format(
                    data_base=self._src.database,
                    table_name=table,
                    exception=','.join(str(i) for i in e.args),
                    select_sql=select_sql
                ))
            await dst_con.commit()
        logger.info('「{}」同步完成，耗时{}s'.format(table, round(time.time() - t, 2)))

    async def main_coroutine(self, name, queue: asyncio.Queue):
        asyncio.current_task().__setattr__('name', name)
        src_con, src_cursor = await self._src.get_db()
        dst_con, dst_cursor = await self._dst.get_db()
        while 1:
            table = await queue.get()
            asyncio.current_task().__setattr__('current_table', table)
            if table == self.SignalQuit:
                logger.debug('coroutine-「{}」结束'.format(name))
                break
            await self._sync_one_table(
                src_con, src_cursor,
                dst_con, dst_cursor,
                table
            )
        src_con.close()
        dst_con.close()

    async def _test_check_coroutine(self, queue: asyncio.Queue):
        while True:
            await asyncio.sleep(5)
            all_tasks_status = []
            for task in asyncio.all_tasks():
                attr = getattr(task, 'current_table', None)
                if attr:
                    all_tasks_status.append(attr)
            if not all_tasks_status:
                break
            logger.debug('剩余任务数量:{}   正在执行:{}'.format(queue.qsize(), all_tasks_status))

    async def pre_scheduler(self):
        await asyncio.gather(self.sync_table_structure())
        if self._err_sql:
            file_name = "err_sql.sql"
            with open(file_name, 'w', encoding="utf-8") as f:
                f.write("\n\n".join(self._err_sql))
            check = input("请确认手动执行{}中的SQL语句后，输入YES执行数据同步:".format(file_name))
            if check != 'YES':
                raise Exception("未确认,数据表同步失败")

    async def main_scheduler(self):
        queue = asyncio.Queue()
        [await queue.put(table) for table in self._sync_tables]
        [await queue.put(self.SignalQuit) for _ in range(self._coroutine)]
        coroutines = [self.main_coroutine('c_{}'.format(i + 1), queue) for i in range(self._coroutine)]
        await asyncio.gather(*coroutines, self._test_check_coroutine(queue))

    def run(self):
        t = time.time()
        # 先执行表结构创建,再同步数据
        asyncio.run(self.pre_scheduler())
        asyncio.run(self.main_scheduler())
        e_t = round(time.time() - t, 2)
        logger.info('同步完成,共计耗时 {}s'.format(e_t))

# todo 索引长度限制导致表创建失败
# todo 数据库连接导致同步进度阻塞,目前只能重新同步
# todo 同步会出现卡主的问题
